# coding:utf-8

"""Ner-trainer&predict

Author:
    name: reeseimk
    email: reeseimk@163.com

Homepage: https://gitee.com/reeseimk/mindspore_bert
"""
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname("__file__"), "./")))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname("__file__"), "../")))
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
import mindspore.numpy as mnp
from mindspore.common.initializer import initializer, Uniform
from mindspore.common.initializer import TruncatedNormal
from mindspore.train.serialization import load_checkpoint, load_param_into_net
import mindspore.dataset as ds
from tqdm import tqdm
from mindspore import context
context.set_context(device_target='CPU')

from src.bert_model import BertModel
from utils.set_config import SetConfig
from utils.deal_data import DealDataNER

def compute_score(emissions, tags, seq_ends, mask, trans, start_trans, end_trans):
    # emissions: (seq_length, batch_size, num_tags)
    # tags: (seq_length, batch_size)
    # mask: (seq_length, batch_size)

    seq_length, batch_size = tags.shape
    mask = mask.astype(emissions.dtype)

    # 将score设置为初始转移概率
    # shape: (batch_size,)
    score = start_trans[tags[0]]
    # score += 第一次发射概率
    # shape: (batch_size,)
    score += emissions[0, mnp.arange(batch_size), tags[0]]

    for i in range(1, seq_length):
        # 标签由i-1转移至i的转移概率（当mask == 1时有效）
        # shape: (batch_size,)
        score += trans[tags[i - 1], tags[i]] * mask[i]

        # 预测tags[i]的发射概率（当mask == 1时有效）
        # shape: (batch_size,)
        score += emissions[i, mnp.arange(batch_size), tags[i]] * mask[i]

    # 结束转移
    # shape: (batch_size,)
    last_tags = tags[seq_ends, mnp.arange(batch_size)]
    # score += 结束转移概率
    # shape: (batch_size,)
    score += end_trans[last_tags]

    return score

def compute_normalizer(emissions, mask, trans, start_trans, end_trans):
    # emissions: (seq_length, batch_size, num_tags)
    # mask: (seq_length, batch_size)

    seq_length = emissions.shape[0]

    # 将score设置为初始转移概率，并加上第一次发射概率
    # shape: (batch_size, num_tags)
    score = start_trans + emissions[0]

    for i in range(1, seq_length):
        # 扩展score的维度用于总score的计算
        # shape: (batch_size, num_tags, 1)
        broadcast_score = score.expand_dims(2)

        # 扩展emission的维度用于总score的计算
        # shape: (batch_size, 1, num_tags)
        broadcast_emissions = emissions[i].expand_dims(1)

        # 根据公式(7)，计算score_i
        # 此时broadcast_score是由第0个到当前Token所有可能路径
        # 对应score的log_sum_exp
        # shape: (batch_size, num_tags, num_tags)
        next_score = broadcast_score + trans + broadcast_emissions

        # 对score_i做log_sum_exp运算，用于下一个Token的score计算
        # shape: (batch_size, num_tags)
        next_score = ops.logsumexp(next_score, axis=1)

        # 当mask == 1时，score才会变化
        # shape: (batch_size, num_tags)
        score = mnp.where(mask[i].expand_dims(1), next_score, score)

    # 最后加结束转移概率
    # shape: (batch_size, num_tags)
    score += end_trans
    # 对所有可能的路径得分求log_sum_exp
    # shape: (batch_size,)
    return ops.logsumexp(score, axis=1)

def viterbi_decode(emissions, mask, trans, start_trans, end_trans):
    # emissions: (seq_length, batch_size, num_tags)
    # mask: (seq_length, batch_size)

    seq_length = mask.shape[0]

    score = start_trans + emissions[0]
    history = ()

    for i in range(1, seq_length):
        broadcast_score = score.expand_dims(2)
        broadcast_emission = emissions[i].expand_dims(1)
        next_score = broadcast_score + trans + broadcast_emission

        # 求当前Token对应score取值最大的标签，并保存
        indices = next_score.argmax(axis=1)
        history += (indices,)

        next_score = next_score.max(axis=1)
        score = mnp.where(mask[i].expand_dims(1), next_score, score)

    score += end_trans

    return score, history

def post_decode(score, history, seq_length):
    # 使用Score和History计算最佳预测序列
    batch_size = seq_length.shape[0]
    seq_ends = seq_length - 1
    # shape: (batch_size,)
    best_tags_list = []

    # 依次对一个Batch中每个样例进行解码
    for idx in range(batch_size):
        # 查找使最后一个Token对应的预测概率最大的标签，
        # 并将其添加至最佳预测序列存储的列表中
        best_last_tag = score[idx].argmax(axis=0)
        best_tags = [int(best_last_tag.asnumpy())]

        # 重复查找每个Token对应的预测概率最大的标签，加入列表
        for hist in reversed(history[:seq_ends[idx]]):
            best_last_tag = hist[idx][best_tags[-1]]
            best_tags.append(int(best_last_tag.asnumpy()))

        # 将逆序求解的序列标签重置为正序
        best_tags.reverse()
        best_tags_list.append(best_tags)

    return best_tags_list

def sequence_mask(seq_length, max_length, batch_first=False):
    """根据序列实际长度和最大长度生成mask矩阵"""
    range_vector = mnp.arange(0, max_length, 1, seq_length.dtype)
    result = range_vector < seq_length.view(seq_length.shape + (1,))
    if batch_first:
        return result.astype(ms.int64)
    return result.astype(ms.int64).swapaxes(0, 1)

class CRF(nn.Cell):
    def __init__(self, num_tags: int, batch_first: bool = False, reduction: str = 'sum') -> None:
        if num_tags <= 0:
            raise ValueError(f'invalid number of tags: {num_tags}')
        super().__init__()
        if reduction not in ('none', 'sum', 'mean', 'token_mean'):
            raise ValueError(f'invalid reduction: {reduction}')
        self.num_tags = num_tags
        self.batch_first = batch_first
        self.reduction = reduction
        self.start_transitions = ms.Parameter(initializer(Uniform(0.1), (num_tags,)), name='start_transitions')
        self.end_transitions = ms.Parameter(initializer(Uniform(0.1), (num_tags,)), name='end_transitions')
        self.transitions = ms.Parameter(initializer(Uniform(0.1), (num_tags, num_tags)), name='transitions')

    def construct(self, emissions, tags=None, seq_length=None):
        if tags is None:
            return self._decode(emissions, seq_length)
        return self._forward(emissions, tags, seq_length)

    def _forward(self, emissions, tags=None, seq_length=None):
        if self.batch_first:
            batch_size, max_length = tags.shape
            emissions = emissions.swapaxes(0, 1)
            tags = tags.swapaxes(0, 1)
        else:
            max_length, batch_size = tags.shape

        if seq_length is None:
            seq_length = mnp.full((batch_size,), max_length, ms.int64)

        mask = sequence_mask(seq_length, max_length)

        # shape: (batch_size,)
        numerator = compute_score(emissions, tags, seq_length-1, mask, self.transitions, self.start_transitions, self.end_transitions)
        # shape: (batch_size,)
        denominator = compute_normalizer(emissions, mask, self.transitions, self.start_transitions, self.end_transitions)
        # shape: (batch_size,)
        llh = denominator - numerator

        if self.reduction == 'none':
            return llh
        if self.reduction == 'sum':
            return llh.sum()
        if self.reduction == 'mean':
            return llh.mean()
        return llh.sum() / mask.astype(emissions.dtype).sum()

    def _decode(self, emissions, seq_length=None):
        if self.batch_first:
            batch_size, max_length = emissions.shape[:2]
            emissions = emissions.swapaxes(0, 1)
        else:
            batch_size, max_length = emissions.shape[:2]

        if seq_length is None:
            seq_length = mnp.full((batch_size,), max_length, ms.int64)

        mask = sequence_mask(seq_length, max_length)

        return viterbi_decode(emissions, mask, self.transitions, self.start_transitions, self.end_transitions)

    
class BiLSTM_CRF(nn.Cell):
    def __init__(self, args, config, is_training, num_tags, dropout_prob=0.0, use_one_hot_embeddings=False):
        super().__init__()
        # self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=padding_idx)
        
        self.currentDir = os.path.abspath(os.path.join(os.path.dirname("__file__"), "./"))
        self.parentDir = os.path.abspath(os.path.join(os.path.dirname("__file__"), "../"))
        
        # if not is_training:
        #     config.hidden_dropout_prob = 0.0
        #     config.hidden_probs_dropout_prob = 0.0
        
        # 构建模型 / 加载bert预训练模型=========================
        # 构建模型
        self.bert = BertModel(config, is_training, use_one_hot_embeddings)
        # self.bert.set_train(True)
        self._param_dict = load_checkpoint(args.pretrain_load_ckpt_path)
        # 预训练权重中的层名称，需要根据模型构建情况进行修改
        param_dict = {}
        for k, v in self._param_dict.items():
            new_k = k.replace('bert.bert.', 'bert.')
            param_dict[new_k] = v
        # 加载权重
        load_param_into_net(self.bert, param_dict)
        # ====================================================
        self.dtype = config.dtype
        self.weight_init = TruncatedNormal(config.initializer_range)
        self.lstm = nn.LSTM(768, 128, num_layers=1, bidirectional=True, batch_first=True)
        self.dense_1 = nn.Dense(256, num_tags, weight_init=self.weight_init, has_bias=True).to_float(config.compute_type)
        # self.hidden2tag = nn.Dense(256, num_tags, 'he_uniform')
        self.dropout = nn.Dropout(p=dropout_prob)
        self.crf = CRF(num_tags, batch_first=True)

    def construct(self, input_ids, token_type_id, input_mask, seq_length, tags=None):
        sequence_output, pooled_output, _ = self.bert(input_ids, token_type_id, input_mask)
        # sequence_output是原始输出，dim=3，bert后接lstm时，需要用这个输出
        # pooled_output是取了CLS的输出，dim=2，bert后直接Dense时，需要用这个输出
        bert_out = self.cast(sequence_output, self.dtype)
        outputs, _ = self.lstm(bert_out)
        outputs = self.dropout(outputs)
        feats = self.dense_1(outputs)

        crf_outs = self.crf(feats, tags, seq_length.reshape(-1))  # seq_length.reshape(-1)需要变成dim=1
        return crf_outs

def load_data_ner(deal_data_ner: DealDataNER=None, batch_size: int=4):
        dataset = ds.GeneratorDataset(deal_data_ner, column_names=["inputs", "input_masks", "segments", "seq_length", "labels"], shuffle=False)
        dataset = dataset.batch(batch_size, drop_remainder=True)
        return dataset

if __name__ == "__main__":
    set_config = SetConfig()
    args, base_bert_cfg = set_config.get_config()
    data_path = "data/ner_data/ner_test.json"
    chunk_path= "data/ner_data/chunk.json"
    max_seq_length = 128
    lr = 1e-5
    wd = 1e-6
    deal_data_ner = DealDataNER(data_path=data_path, chunk_path=chunk_path, max_seq_length=max_seq_length)

    tag_to_idx = deal_data_ner.label_map

    dataset_ner = load_data_ner(deal_data_ner=deal_data_ner, batch_size=4)

    is_training = True
    model = BiLSTM_CRF(args, base_bert_cfg, is_training, len(tag_to_idx))
    optimizer = nn.Adam(model.trainable_params(), learning_rate=lr, weight_decay=wd)

    grad_fn = ms.value_and_grad(model, None, optimizer.parameters)

    def train_step(input_ids, input_mask, segment_ids, seq_length, label_ids):
        loss, grads = grad_fn(input_ids, input_mask, segment_ids, seq_length, label_ids)
        optimizer(grads)
        return loss

    def train_one_epoch(model, dataset, epoch=0):
        model.set_train(True)
        total = dataset.get_dataset_size()
        # print(total)
        loss_total = 0
        step_total = 0
        
        with tqdm(total=total) as t:
            t.set_description('Epoch %i' % epoch)
            for item in dataset.create_tuple_iterator():
                # print(item)
                input_ids, input_mask, segment_ids, seq_length, label_ids = item
                loss = train_step(input_ids, input_mask, segment_ids, seq_length, label_ids)
                loss_total += loss.asnumpy()
                step_total += 1
                loss_mean = loss_total/step_total
                t.set_postfix(loss=loss_mean)
                t.update(1)
        return loss_mean

    epoch_num = 20
    # 训练过程
    for epoch in range(epoch_num): 
        train_loss = train_one_epoch(model, dataset_ner, epoch=epoch)

    # 预测数据准备
    valid_data = {"text": "清华大学坐落于首都北京", "label": [[0,4,"地名"],[9,2,"地名"]]}
    input_ids, input_mask, segment_ids, seq_length, label_ids = deal_data_ner.convert_single_example(valid_data['text'], max_seq_length, valid_data['label'])
    # 需要对input_ids, input_mask, segment_ids升一维到dim=2
    input_ids, input_mask, segment_ids = ms.Tensor([input_ids.asnumpy()], dtype=ms.int32), ms.Tensor([input_mask.asnumpy()], dtype=ms.int32), ms.Tensor([segment_ids.asnumpy()], dtype=ms.int32)
    # print(input_ids, input_mask, segment_ids, seq_length, label_ids)

    # 预测过程
    score, history = model(input_ids, segment_ids, input_mask, seq_length)
    print(score)

    predict = post_decode(score, history, seq_length)
    print(predict)

    idx_to_tag = {idx: tag for tag, idx in tag_to_idx.items()}

    def sequence_to_tag(sequences, idx_to_tag):
        outputs = []
        for seq in sequences:
            outputs.append([idx_to_tag[i] for i in seq])
        return outputs

    res = sequence_to_tag(predict, idx_to_tag)
    print(res)
